import jsonlines
import pickle
import logging
import os
import re
from typing import Dict, List, Any, Tuple

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class DomainIndex:
    def __init__(self):
        """初始化领域分类索引系统"""
        self.domain_index: Dict[str, List[Any]] = {}  # 存储子领域到知识单元的映射
        self.error_records: List[Tuple[int, Any, str]] = []  # 存储错误记录（序号、内容、错误原因）

    def _standardize_domain_name(self, domain: str) -> str:
        """
        标准化领域名称，统一格式
        :param domain: 原始领域名称
        :return: 标准化后的领域名称
        """
        if not domain:
            return "未分类"

        # 去除多余空格和特殊字符
        domain = domain.strip()
        domain = re.sub(r'[^\w\/\u4e00-\u9fa5]', '', domain)  # 保留字母、数字、中文和斜杠
        domain = re.sub(r'\/+', '/', domain)  # 合并多个斜杠
        domain = domain.lower()  # 转为小写，统一大小写

        return domain

    def build_from_jsonlines(self, structured_file: str, index_output_path: str,
                             domain_key: str = "sub_category", skip_errors: bool = False) -> Dict[str, List[Any]]:
        """
        从JSON Lines文件构建领域分类索引
        :param structured_file: 结构化知识单元JSON Lines文件
        :param index_output_path: 索引输出路径
        :param domain_key: 用于分组的子领域字段名
        :param skip_errors: 是否跳过错误记录（True=跳过，False=终止）
        :return: 领域索引字典（key: 子领域，value: 知识单元列表）
        """
        # 检查输入文件是否存在
        if not os.path.exists(structured_file):
            raise FileNotFoundError(f"找不到输入文件: {structured_file}")

        # 重置索引和错误记录
        self.domain_index = {}
        self.error_records = []

        # 按子领域分组知识单元
        logging.info(f"从 {structured_file} 加载知识单元并按子领域分组...")

        with jsonlines.open(structured_file, mode="r") as f:
            total_records = 0
            valid_records = 0

            for i, unit in enumerate(f, 1):
                total_records += 1
                try:
                    # 获取并标准化子领域名称
                    domain = self._standardize_domain_name(unit[domain_key])

                    # 添加到对应领域分组
                    if domain not in self.domain_index:
                        self.domain_index[domain] = []
                    self.domain_index[domain].append(unit)
                    valid_records += 1

                except KeyError as e:
                    error_msg = f"缺少必要的领域字段: {e}"
                    self.error_records.append((i, unit, error_msg))
                    logging.error(f"第{i}条记录{error_msg}")
                    if not skip_errors:
                        raise  # 终止执行
                except Exception as e:
                    error_msg = f"处理出错: {str(e)}"
                    self.error_records.append((i, unit, error_msg))
                    logging.error(f"第{i}条记录{error_msg}")
                    if not skip_errors:
                        raise  # 终止执行

        # 输出统计信息
        logging.info(f"数据处理完成 - 总记录数: {total_records}, 有效记录数: {valid_records}, "
                     f"错误记录数: {len(self.error_records)}")

        # 输出领域分布统计
        self._log_domain_distribution()

        # 确保输出目录存在
        output_dir = os.path.dirname(index_output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)

        # 保存领域索引
        with open(index_output_path, "wb") as f:
            pickle.dump(self.domain_index, f)

        logging.info(f"领域索引构建完成！包含 {len(self.domain_index)} 个子领域，保存至 {index_output_path}")

        # 如果有错误记录，提示用户
        if self.error_records:
            logging.warning(f"注意：共发现 {len(self.error_records)} 条错误记录，已跳过")

        return self.domain_index

    def _log_domain_distribution(self):
        """输出各领域的知识单元数量分布（按数量排序）"""
        if not self.domain_index:
            return

        # 按知识单元数量排序
        sorted_domains = sorted(
            self.domain_index.items(),
            key=lambda x: len(x[1]),
            reverse=True
        )

        logging.info("\n领域分布统计（按知识单元数量排序）:")
        for domain, units in sorted_domains[:5]:  # 显示前5个领域
            logging.info(f"  {domain}: {len(units)} 条")
        if len(sorted_domains) > 5:
            logging.info(f"  ... 还有 {len(sorted_domains) - 5} 个子领域")

    def load_index(self, index_path: str) -> Dict[str, List[Any]]:
        """加载已保存的领域索引"""
        if not os.path.exists(index_path):
            raise FileNotFoundError(f"找不到领域索引文件: {index_path}")

        logging.info(f"从 {index_path} 加载领域索引...")
        with open(index_path, "rb") as f:
            self.domain_index = pickle.load(f)

        logging.info(f"领域索引加载完成，包含 {len(self.domain_index)} 个子领域")
        self._log_domain_distribution()  # 加载时也显示领域分布
        return self.domain_index

    def get_domain_units(self, domain: str, fuzzy: bool = False) -> List[Any]:
        """
        获取指定子领域的所有知识单元
        :param domain: 子领域名称
        :param fuzzy: 是否模糊匹配（包含关键词即可）
        :return: 知识单元列表（不存在则返回空列表）
        """
        if not fuzzy:
            # 精确匹配（标准化后）
            standardized_domain = self._standardize_domain_name(domain)
            return self.domain_index.get(standardized_domain, [])
        else:
            # 模糊匹配（包含关键词）
            matched_units = []
            domain_keyword = domain.lower()
            for domain_name, units in self.domain_index.items():
                if domain_keyword in domain_name:
                    matched_units.extend(units)
            return matched_units

    def list_domains(self, sort_by_count: bool = True) -> List[str]:
        """
        获取所有子领域名称列表
        :param sort_by_count: 是否按知识单元数量排序
        :return: 子领域名称列表
        """
        if not sort_by_count:
            return list(self.domain_index.keys())

        # 按知识单元数量排序
        return [
            domain for domain, _ in sorted(
                self.domain_index.items(),
                key=lambda x: len(x[1]),
                reverse=True
            )
        ]

    def get_error_records(self) -> List[Tuple[int, Any, str]]:
        """获取构建索引过程中的错误记录"""
        return self.error_records


# 使用示例
if __name__ == "__main__":
    # 获取项目根目录
    def get_project_root() -> str:
        current_dir = os.path.dirname(os.path.abspath(__file__))
        project_root = os.path.dirname(current_dir)

        # 验证项目根目录（检查data目录是否存在）
        test_data_dir = os.path.join(project_root, "data")
        if os.path.exists(test_data_dir):
            return project_root
        return current_dir


    # 初始化领域索引
    domain_indexer = DomainIndex()

    # 配置文件路径
    project_root = get_project_root()
    data_dir = os.path.join(project_root, "data")
    jsonlines_file = os.path.join(data_dir, "knowledge_units.jsonl")
    domain_index_path = os.path.join(data_dir, "domain_index.pkl")

    # 确保data目录存在
    os.makedirs(data_dir, exist_ok=True)

    # 构建或加载领域索引（设置skip_errors=True跳过错误记录）
    if not os.path.exists(domain_index_path):
        domain_indexer.build_from_jsonlines(
            jsonlines_file,
            domain_index_path,
            domain_key="sub_category",
            skip_errors=True  # 跳过错误记录，继续处理
        )
    else:
        domain_indexer.load_index(domain_index_path)

    # 示例1：查看所有子领域（按知识单元数量排序）
    print("\n所有子领域列表（按知识单元数量排序）:")
    domains = domain_indexer.list_domains()
    for i, domain in enumerate(domains[:10], 1):
        print(f"{i}. {domain} ({len(domain_indexer.get_domain_units(domain))}条)")
    if len(domains) > 10:
        print(f"... 共 {len(domains)} 个子领域")

    # 示例2：精确查询特定领域
    target_domain = "基站/射频设备参数"
    domain_units = domain_indexer.get_domain_units(target_domain)
    print(f"\n'{target_domain}'领域的知识单元数量: {len(domain_units)}")

    # 示例3：模糊查询领域（包含"基站"关键词）
    fuzzy_results = domain_indexer.get_domain_units("基站", fuzzy=True)
    print(f"\n包含'基站'关键词的所有领域知识单元总数: {len(fuzzy_results)}")

    # 示例4：查看错误记录（如果有）
    errors = domain_indexer.get_error_records()
    if errors:
        print(f"\n错误记录数量: {len(errors)}")
        print(f"第一条错误记录: 第{errors[0][0]}条 - {errors[0][2]}")

    # 示例5：查看领域第一条知识单元
    if domain_units:
        print(f"\n'{target_domain}'第一条知识单元示例: {domain_units[0]['unit_id']}")
        print(f"核心内容: {domain_units[0]['core_content'][:100]}...")
